package com.example.core.filter.flowctl.limiter;

import com.alibaba.fastjson.JSON;
import com.example.common.config.Rule;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import static com.example.common.constants.FilterConst.FLOW_CTL_LIMIT_DURATION;
import static com.example.common.constants.FilterConst.FLOW_CTL_LIMIT_PERMITS;

/**
 * @Author Peng Yisheng
 * @Date 2024/12/7 8:30
 * @Description 使用guava实现单机版限流：基于令牌桶
 */
@Slf4j
public class GuavaCountLimiter {

    private RateLimiter rateLimiter;  // guava限流

    private double maxPermits;  // 每秒最大请求数量

    /**
     * 没有预热，尽可能按照这个速率来分发许可。速率限制器不会考虑之前的请求，也不会允许短时间内的请求速率超过指定的速率。
     * 这可能导致一些请求需要等待，以便速率限制器可以分发足够的许可。
     *
     * @param maxPermits
     */
    public GuavaCountLimiter(double maxPermits) {
        this.maxPermits = maxPermits;
        this.rateLimiter = RateLimiter.create(maxPermits);
    }

    /**
     * 创建一个具有预热期的速率限制器。预热期是指在速率限制器刚开始使用时，速率限制器允许请求超过其平均速率。
     * 预热结束后按照指定的速度分发许可。
     * 提供预热意味着可以在应用启动或负载增加时，允许一些瞬时的高请求速率，然后逐渐调整到稳定的速率。
     *
     * @param maxPermits           表示每秒的最大许可数量
     * @param warmUpPeriodAsSecond 预热时长
     */
    public GuavaCountLimiter(double maxPermits, long warmUpPeriodAsSecond) {
        this.maxPermits = maxPermits;
        this.rateLimiter = RateLimiter.create(maxPermits, warmUpPeriodAsSecond, TimeUnit.SECONDS);
    }

    public static ConcurrentHashMap<String, GuavaCountLimiter> guavaRateLimiterMap = new ConcurrentHashMap<>();

    public static GuavaCountLimiter getInstance(String serviceId, Rule.FlowControlConfig flowControlConfig) {
        if (StringUtils.isEmpty(serviceId) || flowControlConfig == null || StringUtils.isEmpty(flowControlConfig.getValue()) || StringUtils.isEmpty(flowControlConfig.getConfig()) || StringUtils.isEmpty(flowControlConfig.getType())) {
            return null;
        }
        // TODO: 这里没有考虑限流配置变更时缓存需要变更，后续修复
        String key = String.format("%s.%s", serviceId, flowControlConfig.getValue());
        GuavaCountLimiter countLimiter = guavaRateLimiterMap.get(key);
        if (countLimiter == null) {
            // 获得当前路径对应的流控次数
            Map<String, Integer> configMap = JSON.parseObject(flowControlConfig.getConfig(), Map.class);
            // 判断是否包含流控规则
            if (!configMap.containsKey(FLOW_CTL_LIMIT_DURATION) || !configMap.containsKey(FLOW_CTL_LIMIT_PERMITS)) {
                return null;
            }
            // 得到流控时间和时间内限制次数
            double duration = configMap.get(FLOW_CTL_LIMIT_DURATION);
            double permits = configMap.get(FLOW_CTL_LIMIT_PERMITS);
            // 每秒最大请求数
            double maxPermits = permits / duration;
            countLimiter = new GuavaCountLimiter(maxPermits);
            log.info("guava limiter created: key:{}, maxPermits:{}", key, maxPermits);
            guavaRateLimiterMap.putIfAbsent(key, countLimiter);
        }
        return countLimiter;
    }

    public boolean acquire(int permits) {
        return rateLimiter.tryAcquire(permits);
    }

    public boolean acquire() {
        return rateLimiter.tryAcquire();
    }
}
